package postgres

import (
	"context"
	"database/sql"

	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	log "manager/collector/logger"
	"manager/config"
)

var db *gorm.DB

const DBCONTEXTKEY = "DBKEY"

func GetDB(ctx context.Context) *gorm.DB {
	tx, ok := ctx.Value(DBCONTEXTKEY).(*gorm.DB)
	if ok {
		return tx
	}
	return db.WithContext(ctx)
}

func Init(connectAddr string) func() {
	var (
		err   error
		sqlDB *sql.DB
	)

	//connectAddr = "host=134.175.211.197 user=admin password=ebike2023 dbname=manager port=7432 sslmode=disable TimeZone=Asia/Shanghai"
	if db, err = gorm.Open(postgres.Open(connectAddr), &gorm.Config{}); err != nil {
		log.Fatal(err)
	}
	if sqlDB, err = db.DB(); err != nil {
		log.Fatal(err)
	}
	sqlDB.SetMaxIdleConns(50)
	sqlDB.SetMaxOpenConns(500)

	if err = sqlDB.Ping(); err != nil {
		panic(err)
	}

	if config.Conf.Env == "dev" {
		db = db.Debug()
	}

	return func() {
		if err = sqlDB.Close(); err != nil {
			log.Error(err)
		}
	}
}
